from rest_framework import serializers
from django.db.models import Case, When, Func, ExpressionWrapper
from datetime import date, timedelta

from apps.cbsp.models import Cbsp_Bi_Records
from apps.clsbb.models import *
from apps.system.models import Company

from extensions.common.schema import *
from extensions.common.base import *
from extensions.viewsets import *
from extensions.models import *

from apps.statistic.schemas import *


from apps.opoc.models import Opoc_Records
import pandas as pd

import datetime
from datetime import timedelta

now = datetime.datetime.now()
today = now.date()
today_after10=today+timedelta(days=10)

class HomeOverviewViewSet(BaseViewSet, ListModelMixin):
    """首页概览"""

    def get_count(self,data,add_time_begin,add_time_end):

        companies = Company.objects.all()
        res=[]
        for data_item in data:
            counts_all =list(data_item['model'].objects.filter(company__in=companies).values('company','company__id','company__name','company__index').annotate(count=Count('id')).order_by('company__index'))
            counts_yesterday =list(data_item['model'].objects.filter(company__in=companies,add_time__gte=add_time_begin, add_time__lt=add_time_end,).values('company','company__id','company__name','company__index').annotate(count=Count('id')).order_by('company__index'))
            # 遍历所有公司，如果某个公司没有数据，则添加一个 count 为 0 的记录
            for company in companies:
                if not any(count['company__id'] == company.id for count in counts_all):
                    counts_all.append({'company__id': company.id,'company__name': company.name,'company__index': company.index,'count': 0,})
                if not any(count['company__id'] == company.id for count in counts_yesterday):
                    counts_yesterday.append({'company__id': company.id,'company__name': company.name,'company__index': company.index,'count': 0,})
            # 对 counts 中的 count 求和
            count_all_sum = sum(count['count'] for count in counts_all)
            count_yesterday_sum = sum(count['count'] for count in counts_yesterday)
            # 对最终结果按 company__index 排序
            counts_all_sorted = sorted(counts_all, key=lambda count: count['company__index'])
            counts_yesterday_sorted = sorted(counts_yesterday, key=lambda count: count['company__index'])
            # counts_sorted = sorted(counts, key=lambda count: count['count'], reverse=True)
            if counts_all is not None:
                res.append({
                    "counts_all": counts_all_sorted, "count_all_sum": count_all_sum,
                    "counts_yesterday": counts_yesterday_sorted, "count_yesterday_sum": count_yesterday_sum,
                    'type': data_item["type"], 'name': data_item["name"],
                })
        return res

    @extend_schema(responses={200: HomeViewResponse})
    def list(self, request, *args, **kwargs):
        start_time = pendulum.yesterday().to_datetime_string()
        end_time = pendulum.today().to_datetime_string()
        add_time_begin = request.query_params.get('add_time_begin', start_time)
        add_time_end = request.query_params.get('add_time_end', end_time)
        print(add_time_begin,'---',add_time_end)
        data = [
                {'model': Clsbbfak_Records,'type':'fak','name':'方案库'},
                {'model':Clsbbalk_Records,'type':'alk','name':'案例库'},
                {'model':Clsbbtbk_Records,'type':'tbk','name':'投标库'},
                {'model':Clsbbcpk_Records,'type':'cpk','name':'产品库'},
                {'model':Clsbbhzhbk_Records,'type':'hzhbk','name':'合作伙伴库'},
                {'model':Clsbbgszzk_Records,'type':'gszzk','name':'公司资质库'},
                {'model':Clsbbryzzk_Records,'type':'ryzzk','name':'人员资质库'},
                ]

        # # 案例库
        # companies = Company.objects.all()
        # alk_counts = list(Clsbbalk_Records.objects.filter(company__in=companies).values('company','company__id','company__name','company__index').annotate(count=Count('id')).order_by('company__index'))
        # # 遍历所有公司，如果某个公司没有数据，则添加一个 count 为 0 的记录
        # for company in companies:
        #     if not any(count['company__id'] == company.id for count in alk_counts):
        #         alk_counts.append({
        #             'company__id': company.id,
        #             'company__name': company.name,
        #             'company__index': company.index,
        #             'count': 0,
        #         })
        #         # 对 counts 中的 count 求和
        # alk_count_sum = sum(count['count'] for count in alk_counts)
        # # 对最终结果按 company__index 排序
        # # counts_sorted = sorted(alk_counts, key=lambda count: count['company__index'], reverse=True)
        # alk_counts_sorted = sorted(alk_counts, key=lambda count: count['count'], reverse=True)
        # if alk_counts is not None:
        #     # data['alk_count'] = result['alk_count']
        #     data['alk_count'] = alk_counts_sorted
        #     data['alk_count_sum'] = alk_count_sum
        res =self.get_count(data,add_time_begin,add_time_end)
        print(res)
        return Response(data=res, status=status.HTTP_200_OK)

class CompanySerializer(serializers.ModelSerializer):
    class Meta:
        model = Company
        fields = ['id','index', 'name']

class OpocRecordsCountSerializer(serializers.Serializer):
    company = CompanySerializer()
    count = serializers.IntegerField()
    status_1_count = serializers.IntegerField()
    status_2_count = serializers.IntegerField()
    status_minus_1_count = serializers.IntegerField()
    certificate_1_count = serializers.IntegerField()
    certificate_2_count = serializers.IntegerField()
    certificate_3_count = serializers.IntegerField()
    certificate_4_count = serializers.IntegerField()

class OpocOverviewViewSet(BaseViewSet, ListModelMixin):
    """首页概览"""

    @extend_schema(responses={200: HomeViewResponse})
    def list(self, request, *args, **kwargs):
        print(request.user.permissions,request.user.company)
        start_time = pendulum.today().to_datetime_string()
        end_time = pendulum.tomorrow().to_datetime_string()
        data = {
            "opoc_count_all":0,#总人数
            "opoc_count_on":0,#参工
            "opoc_count_waiting":0,#待工
            "opoc_count_quit":0,#离职
            "opoc_count_induction":0,#本月入职
            "opoc_count_climb_high":0,#登高证数
            "opoc_count_electrician":0,#电工证数
            "opoc_count_safety_officer":0,#安全员证数
            "opoc_count_other_person":0,#其它证数
            "opoc_contract_end_count":0,#合同已经过期
            "opoc_contract_end_10days":0,#合同10天内到期
            "opoc_insurance_invalid_count":0,#保险已经过期
            "opoc_insurance_invalid_10days":0,#保险10天内到期
        }
        queryQ = Q(del_flag=False)
        if "opoc_view_all" in request.user.permissions:
            queryQ = Q(del_flag=False)
        else:
            queryQ = queryQ & Q(company=request.user.company)
        queryset = Opoc_Records.objects.filter(queryQ).values()
        opoc_df=pd.DataFrame(queryset)
        # this_month_start = datetime.datetime(now.year, now.month, 1).date()
        this_month_start = datetime.datetime(now.year, now.month, 1).date()

        # this_month_end = (datetime.datetime(now.year, now.month + 1, 1) - timedelta(days=1)).date()
        this_month_end = datetime.datetime(now.year, now.month + 1, 1).date() if now.month < 12 else datetime.datetime(now.year + 1, 1, 1).date()
        print(this_month_start, this_month_end)
        if not opoc_df.empty:
            data['opoc_count_all']=opoc_df.agg('count')['id']
            data['opoc_count_on']=opoc_df[opoc_df.status==1].agg('count')['id']
            data['opoc_count_waiting']=opoc_df[opoc_df.status==2].agg('count')['id']
            data['opoc_count_quit']=opoc_df[opoc_df.status==-1].agg('count')['id']
            mask=opoc_df.contract_signing_begin_date.apply(lambda x:x is not None)
            data['opoc_count_induction']=opoc_df[mask][opoc_df[mask].contract_signing_begin_date.apply(lambda x:(x >=this_month_start and x<this_month_end))].agg('count')['id']
            data['opoc_count_climb_high']=opoc_df[opoc_df.certificate_holding.apply(lambda x:(x is not None) and (1 in x))].agg('count')['id']
            data['opoc_count_electrician']=opoc_df[opoc_df.certificate_holding.apply(lambda x:(x is not None) and (2 in x))].agg('count')['id']
            data['opoc_count_safety_officer']=opoc_df[opoc_df.certificate_holding.apply(lambda x:(x is not None) and (3 in x))].agg('count')['id']
            data['opoc_count_other_person']=opoc_df[opoc_df.certificate_holding.apply(lambda x:(x is not None) and (4 in x))].agg('count')['id']
            mask = opoc_df.contract_signing_end_date.apply(lambda x: x is not None)
            data['opoc_contract_end_count']=opoc_df[mask][opoc_df[mask].contract_signing_end_date.apply(lambda x:(x <today))].agg('count')['id']# 合同到期
            data['opoc_contract_end_10days']=opoc_df[mask][opoc_df[mask].contract_signing_end_date.apply(lambda x:x >=today and x <=today_after10)].agg('count')['id']# 合同10天内到期
            mask = opoc_df.insurance_invalid_date.apply(lambda x: x is not None)
            data['opoc_insurance_invalid_count'] = opoc_df[mask][opoc_df[mask].insurance_invalid_date.apply(lambda x: (x <today))].agg('count')['id']  # 保险到期
            data['opoc_insurance_invalid_10days'] = opoc_df[mask][opoc_df[mask].insurance_invalid_date.apply(lambda x: x >=today  and x <=today_after10)].agg('count')['id']  # 保险10天内到期

        return Response(data=data, status=status.HTTP_200_OK)


class OpocBycityViewSet(BaseViewSet, ListModelMixin):
    """首页概览"""
    queryset = Opoc_Records.objects.filter(del_flag=False)
    serializer_class = OpocRecordsCountSerializer

    def list(self, request, *args, **kwargs):
        queryset = self.get_queryset()
        today = date.today()
        month_start = today.replace(day=1)
        month_end = (today.replace(month=today.month + 1) if today.month < 12 else today.replace(year=today.year + 1,
                                                                                                 month=1)).replace(
            day=1) - timedelta(days=1)
        queryset = queryset.annotate(
            certificate_holding_str=Func(F('certificate_holding'), function='CAST',template='%(function)s(%(expressions)s as CHAR)'),
            days_to_contract_end=ExpressionWrapper(F('contract_signing_end_date') - date.today(),output_field=IntegerField()),
            days_to_insurance_end=ExpressionWrapper(F('insurance_invalid_date') - date.today(),output_field=IntegerField())
        ).values('company').annotate(
            count=Count('company'),
            status_1_count=Count(Case(When(status=1, then=1), output_field=IntegerField())),
            status_2_count=Count(Case(When(status=2, then=1), output_field=IntegerField())),
            status_minus_1_count=Count(Case(When(status=-1, then=1), output_field=IntegerField())),
            certificate_1_count=Count(Case(When(certificate_holding_str__contains=1, then=1), output_field=IntegerField())),
            certificate_2_count=Count(Case(When(certificate_holding_str__contains=2, then=1), output_field=IntegerField())),
            certificate_3_count=Count(Case(When(certificate_holding_str__contains=3, then=1), output_field=IntegerField())),
            certificate_4_count=Count(Case(When(certificate_holding_str__contains=4, then=1), output_field=IntegerField())),
            contract_expired_count=Count(Case(When(days_to_contract_end__lt=0, then=1), output_field=IntegerField())),
            contract_expiring_count=Count(Case(When(days_to_contract_end__lt=10, days_to_contract_end__gte=0, then=1),output_field=IntegerField())),
            insurance_expired_count=Count(Case(When(days_to_insurance_end__lt=0, then=1), output_field=IntegerField())),
            insurance_expiring_count=Count(Case(When(days_to_insurance_end__lt=10, days_to_insurance_end__gte=0, then=1),output_field=IntegerField())),
            joined_this_month_count=Count(Case(When(contract_signing_begin_date__gte=month_start, contract_signing_begin_date__lte=month_end, then=1),output_field=IntegerField()))
        ).order_by('company__index')

        results = []
        for record in queryset:
            company = Company.objects.get(pk=record['company'])
            results.append({
                'company': CompanySerializer(company).data,
                'company_name': CompanySerializer(company).data['name'],
                'count': record['count'],
                'status_1_count': record['status_1_count'],
                'status_2_count': record['status_2_count'],
                'status_minus_1_count': record['status_minus_1_count'],
                'joined_this_month_count': record['joined_this_month_count'],
                'certificate_1_count': record['certificate_1_count'],
                'certificate_2_count': record['certificate_2_count'],
                'certificate_3_count': record['certificate_3_count'],
                'certificate_4_count': record['certificate_4_count'],
                'contract_expired_count': record['contract_expired_count'],
                'contract_expiring_count': record['contract_expiring_count'],
                'insurance_expired_count': record['insurance_expired_count'],
                'insurance_expiring_count': record['insurance_expiring_count']

            })
        return Response(data=results, status=status.HTTP_200_OK)


class CbspOverviewViewSet(BaseViewSet, ListModelMixin):
    """首页概览"""

    @extend_schema(responses={200: HomeViewResponse})
    def list(self, request, *args, **kwargs):
        start_time = pendulum.today().to_datetime_string()
        end_time = pendulum.tomorrow().to_datetime_string()
        data = {
            "cbsp_count_all":0,#总人数
            "cbsp_count_on":0,#参工
            "cbsp_count_waiting":0,#待工
            "cbsp_count_quit":0,#离职
            "cbsp_count_induction":0,#本月入职
            "cbsp_count_climb_high":0,#登高证数
            "cbsp_count_electrician":0,#电工证数
            "cbsp_count_safety_officer":0,#安全员证数
            "cbsp_count_other_person":0,#其它证数
            "cbsp_contract_end_count":0,#合同已经过期
            "cbsp_contract_end_10days":0,#合同10天内到期
            "cbsp_insurance_invalid_count":0,#保险已经过期
            "cbsp_insurance_invalid_10days":0,#保险10天内到期
        }
        queryQ = Q(del_flag=False)
        if "cbsp_view_all" in request.user.permissions:
            queryQ = Q(del_flag=False)
        else:
            queryQ = queryQ & Q(company=request.user.company)
        queryset = Cbsp_Bi_Records.objects.filter(queryQ).values()
        cbsp_df=pd.DataFrame(queryset)
        this_month_start = datetime.datetime(now.year, now.month, 1).date()
        this_month_end = (datetime.datetime(now.year, now.month + 1, 1) - timedelta(days=1)).date()
        if not cbsp_df.empty:
            data['cbsp_count_all']=cbsp_df.agg('count')['id']
            data['cbsp_count_on']=cbsp_df[cbsp_df.status==1].agg('count')['id']
            data['cbsp_count_waiting']=cbsp_df[cbsp_df.status==2].agg('count')['id']
            data['cbsp_count_quit']=cbsp_df[cbsp_df.status==-1].agg('count')['id']
            mask=cbsp_df.contract_signing_begin_date.apply(lambda x:x is not None)
            data['cbsp_count_induction']=cbsp_df[mask][cbsp_df[mask].contract_signing_begin_date.apply(lambda x:(x >=this_month_start and x<=this_month_end))].agg('count')['id']
            data['cbsp_count_climb_high']=cbsp_df[cbsp_df.certificate_holding.apply(lambda x:(x is not None) and (1 in x))].agg('count')['id']
            data['cbsp_count_electrician']=cbsp_df[cbsp_df.certificate_holding.apply(lambda x:(x is not None) and (2 in x))].agg('count')['id']
            data['cbsp_count_safety_officer']=cbsp_df[cbsp_df.certificate_holding.apply(lambda x:(x is not None) and (3 in x))].agg('count')['id']
            data['cbsp_count_other_person']=cbsp_df[cbsp_df.certificate_holding.apply(lambda x:(x is not None) and (4 in x))].agg('count')['id']
            mask = cbsp_df.contract_signing_end_date.apply(lambda x: x is not None)
            data['cbsp_contract_end_count']=cbsp_df[mask][cbsp_df[mask].contract_signing_end_date.apply(lambda x:(x <today))].agg('count')['id']# 合同到期
            data['cbsp_contract_end_10days']=cbsp_df[mask][cbsp_df[mask].contract_signing_end_date.apply(lambda x:x >=today and x <=today_after10)].agg('count')['id']# 合同10天内到期
            mask = cbsp_df.insurance_invalid_date.apply(lambda x: x is not None)
            data['cbsp_insurance_invalid_count'] = cbsp_df[mask][cbsp_df[mask].insurance_invalid_date.apply(lambda x: (x <today))].agg('count')['id']  # 保险到期
            data['cbsp_insurance_invalid_10days'] = cbsp_df[mask][cbsp_df[mask].insurance_invalid_date.apply(lambda x: x >=today  and x <=today_after10)].agg('count')['id']  # 保险10天内到期

        return Response(data=data, status=status.HTTP_200_OK)
    

class CbspBycityViewSet(BaseViewSet, ListModelMixin):
    """首页概览"""
    queryset = Cbsp_Bi_Records.objects.filter(del_flag=False)
    serializer_class = OpocRecordsCountSerializer

    def list(self, request, *args, **kwargs):
        queryset = self.get_queryset()
        today = date.today()
        month_start = today.replace(day=1)
        month_end = (today.replace(month=today.month + 1) if today.month < 12 else today.replace(year=today.year + 1,
                                                                                                 month=1)).replace(
            day=1) - timedelta(days=1)
        queryset = queryset.annotate(
            certificate_holding_str=Func(F('certificate_holding'), function='CAST',template='%(function)s(%(expressions)s as CHAR)'),
            days_to_contract_end=ExpressionWrapper(F('contract_signing_end_date') - date.today(),output_field=IntegerField()),
            days_to_insurance_end=ExpressionWrapper(F('insurance_invalid_date') - date.today(),output_field=IntegerField())
        ).values('company').annotate(
            count=Count('company'),
            status_1_count=Count(Case(When(status=1, then=1), output_field=IntegerField())),
            status_2_count=Count(Case(When(status=2, then=1), output_field=IntegerField())),
            status_minus_1_count=Count(Case(When(status=-1, then=1), output_field=IntegerField())),
            certificate_1_count=Count(Case(When(certificate_holding_str__contains=1, then=1), output_field=IntegerField())),
            certificate_2_count=Count(Case(When(certificate_holding_str__contains=2, then=1), output_field=IntegerField())),
            certificate_3_count=Count(Case(When(certificate_holding_str__contains=3, then=1), output_field=IntegerField())),
            certificate_4_count=Count(Case(When(certificate_holding_str__contains=4, then=1), output_field=IntegerField())),
            contract_expired_count=Count(Case(When(days_to_contract_end__lt=0, then=1), output_field=IntegerField())),
            contract_expiring_count=Count(Case(When(days_to_contract_end__lt=10, days_to_contract_end__gte=0, then=1),output_field=IntegerField())),
            insurance_expired_count=Count(Case(When(days_to_insurance_end__lt=0, then=1), output_field=IntegerField())),
            insurance_expiring_count=Count(Case(When(days_to_insurance_end__lt=10, days_to_insurance_end__gte=0, then=1),output_field=IntegerField())),
            joined_this_month_count=Count(Case(When(contract_signing_begin_date__gte=month_start, contract_signing_begin_date__lte=month_end, then=1),output_field=IntegerField()))
        ).order_by('company__index')

        results = []
        for record in queryset:
            company = Company.objects.get(pk=record['company'])
            results.append({
                'company': CompanySerializer(company).data,
                'company_name': CompanySerializer(company).data['name'],
                'count': record['count'],
                'status_1_count': record['status_1_count'],
                'status_2_count': record['status_2_count'],
                'status_minus_1_count': record['status_minus_1_count'],
                'joined_this_month_count': record['joined_this_month_count'],
                'certificate_1_count': record['certificate_1_count'],
                'certificate_2_count': record['certificate_2_count'],
                'certificate_3_count': record['certificate_3_count'],
                'certificate_4_count': record['certificate_4_count'],
                'contract_expired_count': record['contract_expired_count'],
                'contract_expiring_count': record['contract_expiring_count'],
                'insurance_expired_count': record['insurance_expired_count'],
                'insurance_expiring_count': record['insurance_expiring_count']

            })
        return Response(data=results, status=status.HTTP_200_OK)


__all__ = [
    'HomeOverviewViewSet',
    'OpocOverviewViewSet',
    'CbspOverviewViewSet',
    'CbspBycityViewSet'
]
